﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Management;
using System.Runtime.InteropServices;

namespace Heyo.Class.Helper
{
    public class MftScanner
    {
        private const uint GenericRead = 0x80000000;
        private const int FileShareRead = 0x1;
        private const int FileShareWrite = 0x2;
        private const int OpenExisting = 3;
        private const int FileReadAttributes = 0x80;
        private const int FileNameIinformation = 9;
        private const int FileFlagBackupSemantics = 0x2000000;
        private const int FileOpenForBackupIntent = 0x4000;
        private const int FileOpenByFileId = 0x2000;
        private const int FileOpen = 0x1;
        private const int ObjCaseInsensitive = 0x40;
        private const int FsctlEnumUsnData = 0x900b3;
        private static readonly IntPtr _invalidHandleValue = new IntPtr(-1);

        private static IntPtr _mHCj;
        private static IntPtr _mBuffer;
        private static int _mBufferSize;

        private static string _mDriveLetter;

        private static IntPtr OpenVolume(string szDriveLetter)
        {
            IntPtr hCj = default;
            //// volume handle

            _mDriveLetter = szDriveLetter;
            hCj = NativeMethods.CreateFile(@"\\.\" + szDriveLetter, GenericRead, FileShareRead | FileShareWrite,
                IntPtr.Zero, OpenExisting, 0, IntPtr.Zero);

            return hCj;
        }


        private static void Cleanup()
        {
            if (_mHCj != IntPtr.Zero)
            {
                // Close the volume handle.
                NativeMethods.CloseHandle(_mHCj);
                _mHCj = _invalidHandleValue;
            }

            if (_mBuffer != IntPtr.Zero)
            {
                // Free the allocated memory
                Marshal.FreeHGlobal(_mBuffer);
                _mBuffer = IntPtr.Zero;
            }
        }


        public static IEnumerable<string> EnumerateFiles(string szDriveLetter, string keyWord)
        {
            try
            {
                NativeMethods.UsnRecord usnRecord = default;
                NativeMethods.MftEnumData mft = default;
                int dwRetBytes = 0;
                int cb = 0;
                Dictionary<long, FsNode> dicFrnLookup = new Dictionary<long, FsNode>();
                bool bIsFile = false;

                // This shouldn't be called more than once.
                if (_mBuffer.ToInt32() != 0) throw new Exception("invalid buffer");

                // Assign buffer size
                _mBufferSize = 65536;
                //64KB

                // Allocate a buffer to use for reading records.
                _mBuffer = Marshal.AllocHGlobal(_mBufferSize);

                // correct path
                szDriveLetter = szDriveLetter.TrimEnd('\\');

                // Open the volume handle 
                _mHCj = OpenVolume(szDriveLetter);

                // Check if the volume handle is valid.
                if (_mHCj == _invalidHandleValue)
                {
                    string errorMsg = "Couldn't open handle to the volume.";
                    if (!PermissionHelper.IsAdministrator())
                        errorMsg += "Current user is not administrator";

                    throw new Exception(errorMsg);
                }

                mft.StartFileReferenceNumber = 0;
                mft.LowUsn = 0;
                mft.HighUsn = long.MaxValue;

                do
                {
                    if (NativeMethods.DeviceIoControl(_mHCj, FsctlEnumUsnData, ref mft, Marshal.SizeOf(mft), _mBuffer,
                        _mBufferSize, ref dwRetBytes, IntPtr.Zero))
                    {
                        cb = dwRetBytes;
                        // Pointer to the first record
                        IntPtr pUsnRecord = new IntPtr(_mBuffer.ToInt32() + 8);

                        while (dwRetBytes > 8)
                        {
                            // Copy pointer to USN_RECORD structure.
                            usnRecord = (NativeMethods.UsnRecord) Marshal.PtrToStructure(pUsnRecord,
                                usnRecord.GetType());

                            // The filename within the USN_RECORD.
                            string fileName = Marshal.PtrToStringUni(
                                new IntPtr(pUsnRecord.ToInt32() + usnRecord.FileNameOffset),
                                usnRecord.FileNameLength / 2);

                            bIsFile = !usnRecord.FileAttributes.HasFlag(FileAttributes.Directory);
                            dicFrnLookup.Add(usnRecord.FileReferenceNumber,
                                new FsNode(usnRecord.FileReferenceNumber, usnRecord.ParentFileReferenceNumber, fileName,
                                    bIsFile));

                            // Pointer to the next record in the buffer.
                            pUsnRecord = new IntPtr(pUsnRecord.ToInt32() + usnRecord.RecordLength);

                            dwRetBytes -= usnRecord.RecordLength;
                        }

                        // The first 8 bytes is always the start of the next USN.
                        mft.StartFileReferenceNumber = Marshal.ReadInt64(_mBuffer, 0);
                    }
                    else
                    {
                        break; // TODO: might not be correct. Was : Exit Do
                    }
                } while (!(cb <= 8));

                // Resolve all paths for Files
                foreach (FsNode oFsNode in dicFrnLookup.Values.Where(o => o.IsFile && o.FileName.Contains(keyWord)))
                {
                    string sFullPath = oFsNode.FileName;
                    FsNode oParentFsNode = oFsNode;

                    while (dicFrnLookup.TryGetValue(oParentFsNode.ParentFrn, out oParentFsNode))
                        sFullPath = string.Concat(oParentFsNode.FileName, @"\", sFullPath);
                    sFullPath = string.Concat(szDriveLetter, @"\", sFullPath);

                    yield return sFullPath;
                }
            }
            finally
            {
                //// cleanup
                Cleanup();
            }
        }

        public static IEnumerable<string> EnumerateFiles(string keyWord)
        {
            IEnumerable<string> result = new List<string>();
            SelectQuery selectQuery = new SelectQuery("select * from win32_logicaldisk");
            ManagementObjectSearcher searcher = new ManagementObjectSearcher(selectQuery);
            foreach (ManagementObject disk in searcher.Get())
                result = result.Union(EnumerateFiles(disk["Name"].ToString(), keyWord));
            return result;
        }

        public static IEnumerable<string> EnumerateFiles(DriveInfo drive)
        {
            return EnumerateFiles(drive.Name, "");
        }

        private class FsNode
        {
            public readonly string FileName;
            public long Frn;

            public readonly bool IsFile;
            public readonly long ParentFrn;

            public FsNode(long lFrn, long lParentFsn, string sFileName, bool bIsFile)
            {
                Frn = lFrn;
                ParentFrn = lParentFsn;
                FileName = sFileName;
                IsFile = bIsFile;
            }
        }
    }
}